#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Datetime: 2020/8/6 20:10
# @Author  : CHEN Wang
# @Site    :
# @File    : industry_components.py
# @Software: PyCharm

"""
基于成分股计算行业因子指标，可以是整体法或者中位数法，可以基于全部成分股或者部分成分股（龙头股）
"""

import pandas as pd

from quant_researcher.quant.datasource_fetch.common_data_api import t_trade_date
from quant_researcher.quant.datasource_fetch.factor_api import factor_info, factor_exposure_related
from quant_researcher.quant.datasource_fetch.index_api.index_components import \
    get_index_hist_stock_pool, get_index_dynamic_stock_pool
from quant_researcher.quant.datasource_fetch.index_api.index_constant \
    import SW_CODE_LIST, ZZ_CODE_LIST


# 获取股票池中的龙头股
def get_bellwether_stock(begin, end, stockcodes, **kwargs):
    """
    获取指定区间内，在动态成分股中寻找龙头股(龙头股以换手率因子排名前5评定)
    :param str begin: 指定区间的开始
    :param str end: 指定区间的结束
    :param list stockcodes: 获取的股票代码的范围
    :param kwargs: 额外的过滤
        - pd.Dataframe in_out: 动态成分股的信息， 每天的成分股不一样
            stockcode   000001  002142  002807   ...    601997  601998  603323
            tradedate                            ...
            2019-01-02     1.0     1.0     1.0   ...     0       1.0     1.0
            2019-01-03     1.0     1.0     1.0   ...     1.0     0       1.0
            2019-01-04     1.0     1.0     1.0   ...     1.0     1.0     0

    :return is_bellwether: 一个元素全为0或1的DataFrame，0表示非此股票在某一日期非龙头股，1表示属于龙头股
        stockcode   000001  002142  002807   ...    601997  601998  603323
        tradedate                            ...
        2019-01-02     1.0     1.0     1.0   ...     0       1.0     1.0
        2019-01-03     1.0     1.0     1.0   ...     1.0     0       1.0
        2019-01-04     1.0     1.0     1.0   ...     1.0     1.0     0
    """

    # 以30日平均换手率判断龙头股
    # turnover30 = factor_exposure_related.get_stock_factor_exposure('turnover30', stock_pool=stockcodes,
    #                                                                start=begin, end=end)
    turnover30 = factor_exposure_related.get_stock_factor_exposure('turnover30', stockcodes, begin, end)
    turnover30 = turnover30.rename(columns={'end_date': 'tradedate', 'stock_code': 'stockcode'})
    turnover30 = turnover30.set_index(['tradedate', 'stockcode'])['turnover30'].unstack()
    turnover30.fillna(value=0, inplace=True)
    in_out = kwargs.pop('in_out', None)
    if in_out is not None:
        turnover30 = (turnover30 * in_out).fillna(value=0)
    # 换手率前5个为龙头股， todo 后续再优化
    is_bellwether = turnover30[turnover30.rank(axis=1, ascending=False).apply(
        lambda x: x.isin(range(1, 6)), axis=1)].fillna(value=0)
    is_bellwether[is_bellwether != 0] = 1

    return is_bellwether


# 复合行业因子函数
def get_industry_factor(factor, industry_code=None, industry_type='SW', **method_kwargs):
    """
    通过个股因子构建出的行业因子，在应用时，只需进行市值中性化，无需行业中性化

     :param str factor: 需要进行复合的因子在生产库中的英文名称 目前只支持str（单因子）
     :param str industry_code: 如指定特定行业代码，则对某一行业进行复合；否则计算指定industry_type中，每个行业的因子值
     :param str industry_type: 支持中证行业 ‘ZZ’，申万行业'SW'
     :param method_kwargs:
         - str method: str, 行业因子合成方法， 支持平均(equal)、中位数(median)、市值加权(size_weighted)
         - str universal: str，统计范围，支持全部成分股(all)，市值前n个(eg.3)，前10%(eg.30%)，龙头股(bellwether)
         - list of str time_interval: 计算的起始时间和结束时间，如['2020-01-01', '2020-05-01'], 默认为None,
                                     不指定则计算该因子所有历史时间
         - bool external: True or False, 是否使用外部因子数据
         - pd.Dataframe external_data: 如果使用外部因子数据,在此传入因子数据的DataFrame, 需包含tradedate和stockcode, 其中,
         tradedate需为datetime格式, 所需复合的因子数据的列名需和传入的factor参数一致,默认为None

        一个包含ROE指标的数据形状如下：
        —————————————
        stockcode |tradedate | roe
        000001   |2010-01-03| 0.2
       000002   |2010-01-03| 0.15

     :return: pd.DataFrame
    """
    if industry_code is None:  # 计算某一类型行业的每个行业的因子值
        if industry_type == 'SW':
            industry_code_list = SW_CODE_LIST
        elif industry_type == 'ZZ':
            industry_code_list = ZZ_CODE_LIST
        comp_fac = pd.DataFrame()
        for industry_code in industry_code_list:
            print(f'Processing {industry_code}')
            temp_fac = get_industry_factor_corefunc(factor, industry_code, **method_kwargs)
            temp_fac = temp_fac.set_index('tradedate')
            comp_fac = pd.concat([comp_fac, temp_fac], axis=1)
        comp_fac.index.names = ['tradedate']
        comp_fac.reset_index(inplace=True)
    else:  # 计算指定industry行业的因子值
        print(f'Processing {industry_code}')
        comp_fac = get_industry_factor_corefunc(factor, industry_code, **method_kwargs)

    return comp_fac


# 复合行业因子主要功能函数
def get_industry_factor_corefunc(factor, industry_code=None, **method_kwargs):
    """
    通过个股因子构建出的行业因子，在应用时，只需进行市值中性化，无需行业中性化

     :param str factor: 需要进行复合的因子在生产库中的英文名称 目前只支持str（单因子）
     :param str industry_code: 如指定特定行业代码，则对某一行业进行复合；否则计算指定industry_type中，每个行业的因子值
     :param method_kwargs:
         - str method: str, 行业因子合成方法， 支持平均(equal)、中位数(median)、市值加权(size_weighted)，默认市值加权
         - str universal: str，统计范围，支持全部成分股(all)，市值前n个(eg.3)，前10%(eg.30%)，龙头股(bellwether)，默认支持全部成分股
         - list of str time_interval: 计算的起始时间和结束时间，如['2020-01-01', '2020-05-01'], 默认为None,
                                     不指定则计算该因子所有历史时间
         - bool external: True or False, 是否使用外部因子数据, 默认为False, 不使用外部数据
         - pd.Dataframe external_data: 如果使用外部因子数据,在此传入因子数据的DataFrame, 需包含tradedate和stockcode, 其中,
         tradedate需为datetime格式, 所需复合的因子数据的列名需和传入的factor参数一致,默认为None

        一个包含ROE指标的数据形状如下：
        —————————————
        stockcode |tradedate | roe
        000001   |2010-01-03| 0.2
       000002   |2010-01-03| 0.15

     :return: pd.DataFrame
    """

    method = method_kwargs.pop('method', 'size_weighted')
    universal = method_kwargs.pop('universal', 'all')
    external = method_kwargs.pop('external', False)
    external_data = method_kwargs.pop('external_data', None)
    time_interval = method_kwargs.pop('time_interval', None)
    if time_interval is None:
        if not external:
            # 获取所有available的因子的begin和end
            begin, end = factor_info.get_factor_start_end(factor)
        else:
            begin = min(external_data['tradedate'])
            end = max(external_data['tradedate'])
    else:
        begin = time_interval[0]
        end = time_interval[1]

    if industry_code is None:
        raise RuntimeError('Invalid industry code, please check.')

    # 所有回测区间历史成分股的代码
    stockcodes = get_index_hist_stock_pool(industry_code, begin, end)['stock_code'].tolist()
    # 取市值因子
    # log_size = factor_exposure_related.get_stock_factor_exposure('log_size', stock_pool=stockcodes,
    #                                                              start=begin, end=end)
    log_size = factor_exposure_related.get_stock_factor_exposure('log_size', stockcodes, begin, end)
    log_size = log_size.rename(columns={'end_date': 'tradedate', 'stock_code': 'stockcode'})
    log_size = log_size.set_index(['tradedate', 'stockcode'])
    log_size.columns = ['weight_size']

    # 筛选出成分股的因子数据
    if not external:
        # factor_data = factor_exposure_related.get_stock_factor_exposure(factor, stock_pool=stockcodes,
        #                                                                 start=begin, end=end)
        factor_data = factor_exposure_related.get_stock_factor_exposure(factor, stockcodes, begin, end)
        if factor_data is None:
            factor_data = pd.DataFrame(columns=['tradedate', 'stockcode', factor])
        factor_data = factor_data.rename(columns={'end_date': 'tradedate', 'stock_code': 'stockcode'})
        factor_data = factor_data.set_index(['tradedate', 'stockcode'])
        factor_data.dropna(inplace=True)
    else:
        factor_data = external_data.set_index(['tradedate', 'stockcode']).dropna()
    # factor_data = factor_data.reindex(log_size.index)
    # factor_data = factor_data[f'{factor}'].unstack().ffill()
    # 动态成分股
    components = get_index_dynamic_stock_pool(industry_code, begin, end, trading_date=True)
    components = components.rename(columns={'trade_date': 'tradedate', 'stock_code': 'stockcode'})
    components['tradedate'] = pd.to_datetime(components['tradedate']).dt.strftime('%Y-%m-%d')
    tradedate = components.tradedate.unique().tolist()
    components['in'] = 1
    components = components.set_index(['tradedate', 'stockcode'])
    # tradedate = t_trade_date.get_trade_date_by_mkt(start=begin, end=end)['trade_date'].tolist()
    all = factor_data.join(log_size).join(components).loc[tradedate]
    # all = all.loc[tradedate]
    # all = all.dropna(subset=[f'{factor}'])
    all = all[~all.index.duplicated()].sort_index(level=0)
    factor_data = all[f'{factor}'].unstack().ffill()
    # 成分股因子去极值, 上下2.5%
    if not factor_data.empty:
        factor_data = factor_data.clip(lower=factor_data.quantile(0.025, axis=1), upper=factor_data.quantile(0.975, axis=1),
                                axis=0)
    log_size = all['weight_size'].unstack().ffill()
    # in_out为1或0,1表示某日期该股票是这个行业的成分股,0相反
    in_out = all['in'].unstack().fillna(value=0)
    # scope用来控制范围, 默认是全部成分股
    scope = in_out
    # 取出成分股的动态市值信息
    log_size = (log_size * in_out).fillna(value=0)

    # 全部成分股
    if universal == 'all':
        scope = scope
    # 龙头股
    elif universal == 'bellwether':
        is_bellwether = get_bellwether_stock(begin, end, stockcodes, in_out=in_out)
        scope = (scope * is_bellwether).fillna(value=0)
    # 市值前百分比 或 市值前n个
    elif '%' in universal or isinstance(universal, int):
        if '%' in universal:
            top_num = round(float(universal.replace('%', '')) / 100 * log_size.shape[1])
        elif isinstance(universal, int):
            top_num = round(float(universal))
        top = log_size.rank(axis=1, ascending=False).apply(lambda x: x.isin(range(1, top_num + 1)),
                                                           axis=1)
        scope = top.astype('int')

    if method == 'equal':
        weights = 1 / scope.sum(axis=1)
        ind_factor = (factor_data * scope).mul(weights, axis=0).sum(axis=1)
        ind_factor = ind_factor.rename(f'{industry_code}').reset_index()
    elif method == 'size_weighted':
        log_size = log_size * scope
        weights = log_size.div(log_size.sum(axis=1), axis=0)
        ind_factor = (factor_data * scope * weights).sum(axis=1)
        ind_factor = ind_factor.rename(f'{industry_code}').reset_index()
    elif method == 'median':
        ind_factor = (factor_data * scope).median(axis=1)
        ind_factor = ind_factor.rename(f'{industry_code}').reset_index()

    return ind_factor


if __name__ == '__main__':
    # 指定某一因子，某一行业
    ind_fac = get_industry_factor(factor='bias21', industry_code='801780',
                                  time_interval=['2019-01-01', '2020-05-01'],
                                  method='equal', universal='bellwether')

    # 指定某一因子，某一行业类别
    ind_fac2 = get_industry_factor(factor='e_grow_f1', industry_type='ZZ',
                                   time_interval=['2019-01-01', '2020-05-01'],
                                   method='equal', universal='20%')
